"""
@project: vcat
@author: kang 
@github: https://github.com/fsksf 
@since: 2021/12/22 10:11 PM
"""
from collections import defaultdict
from datetime import datetime
from typing import Union, Dict, Tuple
import numpy as np
from vcat.core.obj import BarData
from vcat.core.constant import DataLevel


class ArrayManager(object):
    """
    单标的单周期
    """

    def __init__(self, size, level: DataLevel, interval: int, contract_id: str):
        """Constructor"""
        self.count: int = 0
        self.size: int = size
        self.inited: bool = False
        self.level: DataLevel = level
        self.interval: int = interval
        self.contract_id = contract_id

        self.open_array: np.ndarray = np.zeros(size)
        self.high_array: np.ndarray = np.zeros(size)
        self.low_array: np.ndarray = np.zeros(size)
        self.close_array: np.ndarray = np.zeros(size)
        self.volume_array: np.ndarray = np.zeros(size)
        self.last_bar_dt = datetime(year=2020, month=1, day=1, hour=1, minute=1, second=1)

    def update_bar(self, bar: BarData) -> None:
        """
        Update new bar data into array manager.
        """
        if bar.level != self.level or bar.interval != self.interval or bar.cid != self.contract_id:
            return
        if bar.dt != self.last_bar_dt:
            self.count += 1
            if not self.inited and self.count >= self.size:
                self.inited = True

            self.open_array[:-1] = self.open_array[1:]
            self.high_array[:-1] = self.high_array[1:]
            self.low_array[:-1] = self.low_array[1:]
            self.close_array[:-1] = self.close_array[1:]
            self.volume_array[:-1] = self.volume_array[1:]

        self.open_array[-1] = bar.open_price
        self.high_array[-1] = bar.high_price
        self.low_array[-1] = bar.low_price
        self.close_array[-1] = bar.close_price
        self.volume_array[-1] = bar.volume

    @property
    def open(self) -> np.ndarray:
        """
        Get open price time series.
        """
        return self.open_array

    @property
    def high(self) -> np.ndarray:
        """
        Get high price time series.
        """
        return self.high_array

    @property
    def low(self) -> np.ndarray:
        """
        Get low price time series.
        """
        return self.low_array

    @property
    def close(self) -> np.ndarray:
        """
        Get close price time series.
        """
        return self.close_array

    @property
    def volume(self) -> np.ndarray:
        """
        Get trading volume time series.
        """
        return self.volume_array

    @property
    def open_interest(self) -> np.ndarray:
        """
        Get trading volume time series.
        """
        return self.open_interest_array

    def sma(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result = talib.SMA(self.close, n)
        if array:
            return result
        return result[-1]

    def ema(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Exponential moving average.
        """
        result = talib.EMA(self.close, n)
        if array:
            return result
        return result[-1]

    def kama(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        KAMA.
        """
        result = talib.KAMA(self.close, n)
        if array:
            return result
        return result[-1]

    def wma(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        WMA.
        """
        result = talib.WMA(self.close, n)
        if array:
            return result
        return result[-1]

    def apo(
        self,
        fast_period: int,
        slow_period: int,
        matype: int = 0,
        array: bool = False
    ) -> Union[float, np.ndarray]:
        """
        APO.
        """
        result = talib.APO(self.close, fast_period, slow_period, matype)
        if array:
            return result
        return result[-1]

    def cmo(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        CMO.
        """
        result = talib.CMO(self.close, n)
        if array:
            return result
        return result[-1]

    def mom(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        MOM.
        """
        result = talib.MOM(self.close, n)
        if array:
            return result
        return result[-1]

    def ppo(
        self,
        fast_period: int,
        slow_period: int,
        matype: int = 0,
        array: bool = False
    ) -> Union[float, np.ndarray]:
        """
        PPO.
        """
        result = talib.PPO(self.close, fast_period, slow_period, matype)
        if array:
            return result
        return result[-1]

    def roc(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        ROC.
        """
        result = talib.ROC(self.close, n)
        if array:
            return result
        return result[-1]

    def rocr(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        ROCR.
        """
        result = talib.ROCR(self.close, n)
        if array:
            return result
        return result[-1]

    def rocp(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        ROCP.
        """
        result = talib.ROCP(self.close, n)
        if array:
            return result
        return result[-1]

    def rocr_100(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        ROCR100.
        """
        result = talib.ROCR100(self.close, n)
        if array:
            return result
        return result[-1]

    def trix(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        TRIX.
        """
        result = talib.TRIX(self.close, n)
        if array:
            return result
        return result[-1]

    def std(self, n: int, nbdev: int = 1, array: bool = False) -> Union[float, np.ndarray]:
        """
        Standard deviation.
        """
        result = talib.STDDEV(self.close, n, nbdev)
        if array:
            return result
        return result[-1]

    def obv(self, array: bool = False) -> Union[float, np.ndarray]:
        """
        OBV.
        """
        result = talib.OBV(self.close, self.volume)
        if array:
            return result
        return result[-1]

    def cci(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Commodity Channel Index (CCI).
        """
        result = talib.CCI(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def atr(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Average True Range (ATR).
        """
        result = talib.ATR(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def natr(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        NATR.
        """
        result = talib.NATR(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def rsi(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Relative Strenght Index (RSI).
        """
        result = talib.RSI(self.close, n)
        if array:
            return result
        return result[-1]

    def macd(
        self,
        fast_period: int,
        slow_period: int,
        signal_period: int,
        array: bool = False
    ) -> Union[
        Tuple[np.ndarray, np.ndarray, np.ndarray],
        Tuple[float, float, float]
    ]:
        """
        MACD.
        """
        macd, signal, hist = talib.MACD(
            self.close, fast_period, slow_period, signal_period
        )
        if array:
            return macd, signal, hist
        return macd[-1], signal[-1], hist[-1]

    def adx(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        ADX.
        """
        result = talib.ADX(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def adxr(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        ADXR.
        """
        result = talib.ADXR(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def dx(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        DX.
        """
        result = talib.DX(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def minus_di(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        MINUS_DI.
        """
        result = talib.MINUS_DI(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def plus_di(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        PLUS_DI.
        """
        result = talib.PLUS_DI(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def willr(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        WILLR.
        """
        result = talib.WILLR(self.high, self.low, self.close, n)
        if array:
            return result
        return result[-1]

    def ultosc(
        self,
        time_period1: int = 7,
        time_period2: int = 14,
        time_period3: int = 28,
        array: bool = False
    ) -> Union[float, np.ndarray]:
        """
        Ultimate Oscillator.
        """
        result = talib.ULTOSC(self.high, self.low, self.close, time_period1, time_period2, time_period3)
        if array:
            return result
        return result[-1]

    def trange(self, array: bool = False) -> Union[float, np.ndarray]:
        """
        TRANGE.
        """
        result = talib.TRANGE(self.high, self.low, self.close)
        if array:
            return result
        return result[-1]

    def boll(
        self,
        n: int,
        dev: float,
        array: bool = False
    ) -> Union[
        Tuple[np.ndarray, np.ndarray],
        Tuple[float, float]
    ]:
        """
        Bollinger Channel.
        """
        mid = self.sma(n, array)
        std = self.std(n, 1, array)

        up = mid + std * dev
        down = mid - std * dev

        return up, down

    def keltner(
        self,
        n: int,
        dev: float,
        array: bool = False
    ) -> Union[
        Tuple[np.ndarray, np.ndarray],
        Tuple[float, float]
    ]:
        """
        Keltner Channel.
        """
        mid = self.sma(n, array)
        atr = self.atr(n, array)

        up = mid + atr * dev
        down = mid - atr * dev

        return up, down

    def donchian(
        self, n: int, array: bool = False
    ) -> Union[
        Tuple[np.ndarray, np.ndarray],
        Tuple[float, float]
    ]:
        """
        Donchian Channel.
        """
        up = talib.MAX(self.high, n)
        down = talib.MIN(self.low, n)

        if array:
            return up, down
        return up[-1], down[-1]

    def aroon(
        self,
        n: int,
        array: bool = False
    ) -> Union[
        Tuple[np.ndarray, np.ndarray],
        Tuple[float, float]
    ]:
        """
        Aroon indicator.
        """
        aroon_up, aroon_down = talib.AROON(self.high, self.low, n)

        if array:
            return aroon_up, aroon_down
        return aroon_up[-1], aroon_down[-1]

    def aroonosc(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Aroon Oscillator.
        """
        result = talib.AROONOSC(self.high, self.low, n)

        if array:
            return result
        return result[-1]

    def minus_dm(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        MINUS_DM.
        """
        result = talib.MINUS_DM(self.high, self.low, n)

        if array:
            return result
        return result[-1]

    def plus_dm(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        PLUS_DM.
        """
        result = talib.PLUS_DM(self.high, self.low, n)

        if array:
            return result
        return result[-1]

    def mfi(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Money Flow Index.
        """
        result = talib.MFI(self.high, self.low, self.close, self.volume, n)
        if array:
            return result
        return result[-1]

    def ad(self, array: bool = False) -> Union[float, np.ndarray]:
        """
        AD.
        """
        result = talib.AD(self.high, self.low, self.close, self.volume)
        if array:
            return result
        return result[-1]

    def adosc(
        self,
        fast_period: int,
        slow_period: int,
        array: bool = False
    ) -> Union[float, np.ndarray]:
        """
        ADOSC.
        """
        result = talib.ADOSC(self.high, self.low, self.close, self.volume, fast_period, slow_period)
        if array:
            return result
        return result[-1]

    def bop(self, array: bool = False) -> Union[float, np.ndarray]:
        """
        BOP.
        """
        result = talib.BOP(self.open, self.high, self.low, self.close)

        if array:
            return result
        return result[-1]


class XArrayManager(object):
    """
    同周期多标的行情存储
    """

    def __init__(self, size, level: DataLevel, interval: int):
        """Constructor"""
        self.count: int = 0
        self.size: int = size
        self._inited: bool = False
        self.level: DataLevel = level
        self.interval: int = interval

        self._am_dict: Dict[str, ArrayManager] = {}

    def update_bar(self, bar: BarData) -> None:
        """
        Update new bar data into array manager.
        """
        am = self._am_dict.get(bar.cid)
        if am is None:
            am = ArrayManager(size=self.size, level=self.level, interval=self.interval, contract_id=bar.cid)
            self._am_dict[bar.cid] = am
        am.update_bar(bar)

    @property
    def inited(self):
        if self._inited:
            return True
        for am in self._am_dict.values():
            if not am.inited:
                return False
        self._inited = True
        return self._inited

    def sma(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        Simple moving average.
        """
        return {k: am.sma(n, array) for k, am in self._am_dict.items()}

    def ema(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        Exponential moving average.
        """
        return {k: am.ema(n, array) for k, am in self._am_dict.items()}

    def kama(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        KAMA.
        """
        return {k: am.kama(n, array) for k, am in self._am_dict.items()}

    def wma(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        WMA.
        """
        return {k: am.wma(n, array) for k, am in self._am_dict.items()}

    def apo(
            self,
            fast_period: int,
            slow_period: int,
            matype: int = 0,
            array: bool = False
    ) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        APO.
        """
        return {k: am.apo(fast_period, slow_period, matype, array) for k, am in self._am_dict.items()}

    def cmo(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        CMO.
        """
        return {k: am.cmo(n, array) for k, am in self._am_dict.items()}

    def mom(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        MOM.
        """
        return {k: am.mom(n, array) for k, am in self._am_dict.items()}

    def ppo(
            self,
            fast_period: int,
            slow_period: int,
            matype: int = 0,
            array: bool = False
    ) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        PPO.
        """
        return {k: am.ppo(fast_period, slow_period, matype, array) for k, am in self._am_dict.items()}

    def roc(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        ROC.
        """
        return {k: am.roc(n, array) for k, am in self._am_dict.items()}

    def rocr(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        ROCR.
        """
        return {k: am.rocr(n, array) for k, am in self._am_dict.items()}

    def rocp(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        ROCP.
        """
        return {k: am.rocp(n, array) for k, am in self._am_dict.items()}

    def rocr_100(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        ROCR100.
        """
        return {k: am.rocr_100(n, array) for k, am in self._am_dict.items()}

    def trix(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        TRIX.
        """
        return {k: am.trix(n, array) for k, am in self._am_dict.items()}

    def std(self, n: int, nbdev: int = 1, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        Standard deviation.
        """
        return {k: am.std(n, array) for k, am in self._am_dict.items()}

    def obv(self, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        OBV.
        """
        return {k: am.obv(array) for k, am in self._am_dict.items()}

    def cci(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        Commodity Channel Index (CCI).
        """
        return {k: am.cci(n, array) for k, am in self._am_dict.items()}

    def atr(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        Average True Range (ATR).
        """
        return {k: am.atr(n, array) for k, am in self._am_dict.items()}

    def natr(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        NATR.
        """
        return {k: am.natr(n, array) for k, am in self._am_dict.items()}

    def rsi(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        Relative Strenght Index (RSI).
        """
        return {k: am.rsi(n, array) for k, am in self._am_dict.items()}

    def macd(
            self,
            fast_period: int,
            slow_period: int,
            signal_period: int,
            array: bool = False
    ) -> Union[
        Tuple[np.ndarray, np.ndarray, np.ndarray],
        Tuple[float, float, float]
    ]:
        """
        MACD.
        """
        return {k: am.macd(fast_period, slow_period, signal_period, array) for k, am in self._am_dict.items()}

    def adx(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        ADX.
        """
        return {k: am.adx(n, array) for k, am in self._am_dict.items()}

    def adxr(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        ADXR.
        """
        return {k: am.adxr(n, array) for k, am in self._am_dict.items()}

    def dx(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        DX.
        """
        return {k: am.dx(n, array) for k, am in self._am_dict.items()}

    def minus_di(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        MINUS_DI.
        """
        return {k: am.minus_di(n, array) for k, am in self._am_dict.items()}

    def plus_di(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        PLUS_DI.
        """
        return {k: am.plus_di(n, array) for k, am in self._am_dict.items()}

    def willr(self, n: int, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        WILLR.
        """
        return {k: am.willr(n, array) for k, am in self._am_dict.items()}

    def ultosc(
            self,
            time_period1: int = 7,
            time_period2: int = 14,
            time_period3: int = 28,
            array: bool = False
    ) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        Ultimate Oscillator.
        """
        return {k: am.ultosc(time_period1, time_period2, time_period3, array) for k, am in self._am_dict.items()}

    def trange(self, array: bool = False) -> Union[Dict[str, float], Dict[str, np.ndarray]]:
        """
        TRANGE.
        """
        return {k: am.trange(array) for k, am in self._am_dict.items()}

    def boll(
            self,
            n: int,
            dev: float,
            array: bool = False
    ):
        """
        Bollinger Channel.
        """
        return {k: am.boll(n, dev, array) for k, am in self._am_dict.items()}

    def keltner(
            self,
            n: int,
            dev: float,
            array: bool = False
    ) -> Union[
        Dict[str, Tuple[np.ndarray, np.ndarray]],
        Dict[str, Tuple[float, float]]
    ]:
        """
        Keltner Channel.
        """
        return {k: am.keltner(n, dev, array) for k, am in self._am_dict.items()}

    def donchian(
            self, n: int, array: bool = False
    ) -> Union[Dict[str, Tuple[float, float]], Dict[str, Tuple[np.ndarray, np.ndarray]]]:
        """
        Donchian Channel.
        """
        return {k: am.rocr(n, array) for k, am in self._am_dict.items()}

    def aroon(
            self,
            n: int,
            array: bool = False
    ) -> Union[
        Dict[str, Tuple[np.ndarray, np.ndarray]],
        Dict[str, Tuple[float, float]]
    ]:
        """
        Aroon indicator.
        """
        return {k: am.rocr(n, array) for k, am in self._am_dict.items()}

    def aroonosc(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Aroon Oscillator.
        """
        result = talib.AROONOSC(self.high, self.low, n)

        if array:
            return result
        return result[-1]

    def minus_dm(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        MINUS_DM.
        """
        result = talib.MINUS_DM(self.high, self.low, n)

        if array:
            return result
        return result[-1]

    def plus_dm(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        PLUS_DM.
        """
        result = talib.PLUS_DM(self.high, self.low, n)

        if array:
            return result
        return result[-1]

    def mfi(self, n: int, array: bool = False) -> Union[float, np.ndarray]:
        """
        Money Flow Index.
        """
        result = talib.MFI(self.high, self.low, self.close, self.volume, n)
        if array:
            return result
        return result[-1]

    def ad(self, array: bool = False) -> Union[float, np.ndarray]:
        """
        AD.
        """
        result = talib.AD(self.high, self.low, self.close, self.volume)
        if array:
            return result
        return result[-1]

    def adosc(
            self,
            fast_period: int,
            slow_period: int,
            array: bool = False
    ) -> Union[float, np.ndarray]:
        """
        ADOSC.
        """
        result = talib.ADOSC(self.high, self.low, self.close, self.volume, fast_period, slow_period)
        if array:
            return result
        return result[-1]

    def bop(self, array: bool = False) -> Union[float, np.ndarray]:
        """
        BOP.
        """
        result = talib.BOP(self.open, self.high, self.low, self.close)

        if array:
            return result
        return result[-1]
